# encoding:utf-8
# -----------------------------------------------------------
# "ML-Stealer: Stealing Prediction Functionality of Machine Learning Models with Mere Black-Box Access"
# @author: Shijie Wang, 2019.
# ------------------------------------------------------------

"""
VAE Model
"""
from __future__ import print_function
import argparse
import torch
import random
from torch.utils.data import Dataset, DataLoader
from torch import nn, optim
from torch.nn import functional as F
import torch.nn.parallel
import pandas as pd
import numpy as np


class AdultDataset(Dataset):
    def __init__(self, csv_file):
        self.adult_data = pd.read_csv(csv_file, sep=',')
        # normalize data
        self.adult_data = self.adult_data.apply(lambda x: (x - np.mean(x)) / np.std(x))

    def __len__(self):
        return len(self.adult_data)

    def __getitem__(self, idx):
        return torch.tensor(np.array(self.adult_data.iloc[[idx]])).float().squeeze()


class VAE_adult(nn.Module):
    def __init__(self):
        super(VAE_adult, self).__init__()
        self.fc1 = nn.Linear(14, 256)
        self.fc21 = nn.Linear(256, 3)
        self.fc22 = nn.Linear(256, 3)
        self.fc3 = nn.Linear(3, 256)
        self.fc4 = nn.Linear(256, 14)

    def encode(self, x):
        h1 = F.relu(self.fc1(x))
        return self.fc21(h1), self.fc22(h1)

    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5*logvar)
        eps = torch.randn_like(std)
        return mu + eps*std

    def decode(self, z):
        h3 = F.relu(self.fc3(z))
        return self.fc4(h3)

    def forward(self, x):
        mu, logvar = self.encode(x.view(-1, 14))
        z = self.reparameterize(mu, logvar)
        return self.decode(z), mu, logvar


def loss_function(recon_x, x, mu, logvar):
    MSE = F.mse_loss(recon_x, x.view(-1, 14), reduction='sum')
    KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
    return MSE + KLD


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='VAE data augmentation')
    parser.add_argument('--dataroot', required=True, help='path to dataset')
    parser.add_argument('--workers', type=int, default=2, help='number of data loading workers')
    parser.add_argument('--batchsize', type=int, default=50, help='input batch size')
    parser.add_argument('--epochs', type=int, default=10, help='number of epochs to train')
    parser.add_argument('--cuda', action='store_true', help='enables cuda training')
    parser.add_argument('--ngpu', type=int, default=1, help='number of GPUs to use')
    parser.add_argument('--output', required=True, help='folder to save model parameters')
    parser.add_argument('--manualSeed', type=int, help='manual seed')

    args = parser.parse_args()
    print(args)

    if args.manualSeed is None:
        args.manualSeed = random.randint(1, 10000)
    print("Random Seed: ", args.manualSeed)
    torch.manual_seed(args.manualSeed)
    device = torch.device("cuda" if args.cuda else "cpu")

    model = VAE_adult().to(device)
    optimizer = optim.Adam(model.parameters(), lr=1e-2)

    adult_dataset = AdultDataset(csv_file=args.dataroot)
    adult_dataloader = DataLoader(adult_dataset, batch_size=args.batchsize,
                                  shuffle=True, num_workers=args.workers, pin_memory=True)

    for epoch in range(1, args.epochs + 1):
        train_loss = 0
        mu_mean = torch.zeros(3).to(device)
        logvar_mean = torch.zeros(3).to(device)
        for i, data in enumerate(adult_dataloader):
            data = data.to(device)
            optimizer.zero_grad()
            recon_batch, mu, logvar = model(data)
            loss = loss_function(recon_batch, data, mu, logvar)
            loss.backward()
            train_loss += loss.item()
            mu_mean += torch.mean(mu, 0)
            logvar_mean += torch.mean(logvar, 0)
            optimizer.step()

            if i % 10 == 0:
                print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                    epoch, i * len(data), len(adult_dataloader.dataset),
                    100. * i / len(adult_dataloader),
                    loss.item() / len(data)
                ))

        print('=====>Epoch: {} Average loss: {:.4f}'.format(
            epoch, train_loss / len(adult_dataloader.dataset)
        ))
        print('mu: {}, logvar: {}'.format(mu_mean / len(adult_dataloader), logvar_mean / len(adult_dataloader)))

    torch.save(model, '%s/VAE_adult_svm_no_generator.pth' % args.output)
